#!/usr/bin/env python
# -*- coding: utf-8 -*-
# Copyright 2018 Tongdun Inc. All Rights Reserved.
# @Time             : 2019/2/19 10:15 AM
# @Author           : ding
# @File             : TxtGen.py
# @description      : 创建文本相关资源
import os
import codecs
import random
import time

from PIL import Image, ImageFont, ImageDraw
from tqdm import tqdm

from imprec import voc_keys


def mirror_image_width(image, width):
    w, h = image.size
    if w > width:
        return image, True

    dw = min(width - w, w)

    box = [w - dw, 0, w, h]
    region = image.crop(box)
    region = region.transpose(Image.FLIP_LEFT_RIGHT)

    rw, _ = region.size
    mirror_image = Image.new(image.mode, [w + dw, h])
    mirror_image.paste(image, (0, 0))
    mirror_image.paste(region, (w, 0))

    return mirror_image, w + dw >= width


def mirror_image_height(image, height):
    w, h = image.size
    if h > height:
        return image, True

    dh = min(height - h, h)

    box = [0, h - dh, w, h]
    region = image.crop(box)
    region = region.transpose(Image.FLIP_TOP_BOTTOM)

    mirror_image = Image.new(image.mode, [w, h + dh])
    mirror_image.paste(image, (0, 0))
    mirror_image.paste(region, (0, h))

    return mirror_image, h + dh >= height


def mirror_image(image, width, height):
    while True:
        image, wflag = mirror_image_width(image, width)
        image, hflag = mirror_image_height(image, height)
        if wflag and hflag:
            break
    return image


class TxtImgGen:
    def __init__(self, font_size, back_list=[], font_list=[], back_size=[512, 512], use_bar=True, valid=True):

        self.back_list = []
        if isinstance(back_list, list):
            for blst in back_list:
                with codecs.open(blst, "r", encoding='utf-8') as f:
                    for line in f.readlines():
                        self.back_list.append(line.strip())
        else:
            with codecs.open(back_list, "r", encoding='utf-8') as f:
                for line in f.readlines():
                    self.back_list.append(line.strip())

        self.valid = valid

        self.font_list = []
        if isinstance(font_list, list):
            for flst in font_list:
                with codecs.open(flst, "r", encoding='utf-8') as f:
                    if use_bar:
                        lines = tqdm(f.readlines(), desc="font")
                    else:
                        lines = f.readlines()
                    for line in lines:
                        dict_mask = []

                        if self.valid:
                            with codecs.open(line.strip() + '.lst', "r", encoding='utf-8') as mf:
                                # print(line.strip() + '.lst')
                                for kline in mf.readlines():
                                    try:
                                        key, val = kline.strip().split()
                                    except ValueError:
                                        key = " "
                                        val = kline.strip().split()

                                    if val != 0:
                                        dict_mask.append(key)

                        self.font_list.append([line.strip(), dict_mask])
        else:
            with codecs.open(font_list, "r", encoding='utf-8') as f:
                if use_bar:
                    lines = tqdm(f.readlines(), desc="font")
                else:
                    lines = f.readlines()
                for line in lines:
                    dict_mask = []

                    if self.valid:
                        with codecs.open(line.strip() + '.lst', "r", encoding='utf-8') as mf:
                            # print(line.strip() + '.lst')
                            for kline in mf.readlines():
                                try:
                                    key, val = kline.strip().split()
                                except ValueError:
                                    key = " "
                                    val = kline.strip().split()

                                if val != 0:
                                    dict_mask.append(key)

                    self.font_list.append([line.strip(), dict_mask])

        self.back_size = back_size
        self.font_size = font_size

        self.color = [0, 0, 0]

    def add_back(self, back_list, use_bar=True):
        with codecs.open(back_list, "r", encoding='utf-8') as f:
            if use_bar:
                lines = tqdm(f.readlines(), desc="back")
            else:
                lines = f.readlines()

            for line in lines:
                self.back_list.append(line.strip())

    def add_font(self, font_list, use_bar=True):
        with codecs.open(font_list, "r", encoding='utf-8') as f:
            if use_bar:
                lines = tqdm(f.readlines(), desc="font")
            else:
                lines = f.readlines()
            for line in lines:
                dict_mask = []

                if self.valid:
                    with codecs.open(line.strip() + '.lst', "r", encoding='utf-8') as mf:
                        # print(line.strip() + '.lst')
                        for kline in mf.readlines():
                            try:
                                key, val = kline.strip().split()
                            except ValueError:
                                key = " "
                                val = kline.strip().split()

                            if val != 0:
                                dict_mask.append(key)

                self.font_list.append([line.strip(), dict_mask])

    def __gen_back(self, width, height, enlarge=0):
        """
        产生一个背景
        :param width: 图像的高度
        :param height: 图像的宽度
        :param use_cur: 是否使用上次的图像
        :param enlarge: 默认图像扩展方式，0:resize,other:镜像扩展
        :return:
        """
        try:
            back_path = random.choice(self.back_list)
            img_bkb = Image.open(back_path)
            if img_bkb.mode != "RGBA":
                img_bkb = img_bkb.convert("RGBA")
        except IOError:
            img_bkb = Image.new("RGBA", self.back_size, 255)

        w, h = img_bkb.size
        t = ((height > h) << 1) + (width > w)

        if enlarge == 0 and t != 0:
            return img_bkb.resize([width, height])

        if t == 1:
            # 图像宽度较小，镜像水平扩展
            img_bkb = mirror_image(img_bkb, width, h)
        elif t == 2:
            # 图像高度较小， 镜像垂直扩展
            img_bkb = mirror_image(img_bkb, w, height)
        elif t == 3:
            img_bkb = mirror_image(img_bkb, width, height)
        else:
            pass

        return img_bkb

    def gen_font(self):
        try:
            font, mask = random.choice(self.font_list)
            size = random.choice(self.font_size)
            font = ImageFont.truetype(font, size)
        except IOError:
            font = None

        return mask, font

    def gen_color(self, color=[-1, -1, -1]):
        if isinstance(color[0], list):
            self.color = random.choice(color)
        else:

            if color[0] < 0:
                self.color[0] = random.randint(0, 255)
            else:
                self.color[0] = color[0]

            if color[1] < 0:
                self.color[1] = random.randint(0, 255)
            else:
                self.color[0] = color[1]

            if color[2]:
                self.color[2] = random.randint(0, 255)
            else:
                self.color[2] = color[2]

    def __shack(self, box, size, shack=[[0, 0], [0, 0], [0, 0], [0, 0]]):
        w, h = size
        left = box[0] + random.randint(shack[0][0], shack[0][1])
        left = max(left, 0)
        top = box[1] - random.randint(shack[1][0], shack[1][1])
        top = max(top, 0)
        right = box[2] + random.randint(shack[2][0], shack[2][1])
        right = min(right, w - 1)
        bottom = box[3] + random.randint(shack[3][0], shack[3][1])
        bottom = min(bottom, h - 1)

        return [min(left, right), min(top, bottom), max(left, right), max(top, bottom)]

    def __check_mask(self, text, mask):
        if not self.valid:
            return True

        for ch in text:
            if ch not in mask:
                return False
        return True

    def crop_text_line(self, text, font, mask, opt=[-1, -1], enlarge=-1, shack=[[0, 0], [0, 0], [0, 0], [0, 0]]):
        """
        1. 判断文字是否超出背景
            1. 缩放扩展
            2. 镜像扩展
        2. 判断颜色是否需需要改变
        3. 随机点的位置，并截取图像
        """
        if not self.__check_mask(text, mask):
            return None

        fw, fh = font.getsize(text)

        if enlarge < 0:
            enlarge = random.choice([0, 1])

        image = self.__gen_back(fw, fh, enlarge)
        w, h = image.size

        if opt[0] == -1 or opt[1] == -1:
            x = random.randint(0, w - fw)
            y = random.randint(0, h - fh)
        else:
            x, y = opt

        draw = ImageDraw.Draw(image)
        draw.text((x, y), text=text, fill=tuple(self.color), font=font, align='center')

        box = [x, y, x + fw - 1, y + fh - 1]

        box = self.__shack(box, image.size, shack)

        return image.crop(box)


class TxtGen:
    def __init__(self, list_file=[], use_bar=True, gap=' '):
        self.words = []
        self.use_bar = use_bar
        if isinstance(list_file, list):
            self.read_from_file(list_file, gap)
        else:
            self.read_from_file([list_file], gap)

        self.china_acode = []
        with codecs.open('./china_acode.lst', "r", encoding='utf-8') as f:
            for line in f.readlines():
                self.china_acode.append(line.strip())

    def read_from_file(self, file_list, gap=' '):
        for flst in file_list:
            with codecs.open(flst, "r", encoding='utf-8') as f:
                if self.use_bar:
                    lines = tqdm(f.readlines(), desc=os.path.basename(flst))
                else:
                    lines = f.readlines()
                for line in lines:
                    words = line.strip().split(gap)
                    for word in words:
                        if len(word) > 0:
                            self.words.append(word)

    def gen_text(self, numb_text=1):
        if numb_text <= 1:
            return random.choice(self.words)
        else:
            return random.sample(self.words, numb_text)

    @staticmethod
    def gen_dict_text(numb_words=[1, 5]):
        key = voc_keys.alphabet
        word_of_text = random.randint(numb_words[0], numb_words[1])
        return ''.join(random.sample(key, word_of_text))

    @staticmethod
    def gen_date(format="%Y-%m-%d", start=(1900, 1, 1, 0, 0, 0, 0, 0, 0), stop=(2018, 12, 31, 23, 59, 59, 0, 0, 0)):
        stime = time.mktime(start)
        etime = time.mktime(stop)
        random_time = random.randint(stime, etime)
        date = time.localtime(random_time)

        return time.strftime(format, date)

    def gen_id(self, birthday):
        sex = random.choice([u'男', u'女'])
        chcode = random.choice(self.china_acode)  # [1:6]

        date = birthday  # [7:14]

        two = random.randint(10, 99)  # [15:16]
        if sex == u"女":
            one = random.choice([0, 2, 4, 6, 8])  # [17]
        else:
            one = random.choice([1, 3, 5, 7, 9])  # [17]

        last = random.randint(0, 10)
        if last == 10:
            last = 'X'

        ndate = ""
        for ch in date:
            if ch.isdigit():
                ndate += ch

        id_comb = str(chcode) + ndate + str(two) + str(one) + str(last), str(date), sex
        return id_comb

    @staticmethod
    def gen_nation(numb_text=1):
        nation = [u"汉", u"壮", u"满", u"回", u"苗", u"维吾尔",
                  u"土家", u"彝", u"蒙古", u"藏", u"布依", u"侗", u"瑶",
                  u"朝鲜", u"白", u"哈尼", u"哈萨克", u"黎", u"傣", u"畲",
                  u"傈僳", u"仡佬", u"东乡", u"高山", u"拉祜", u"水", u"佤",
                  u"纳西", u"羌", u"土", u"仫佬", u"锡伯", u"柯尔克孜", u"达斡尔",
                  u"景颇", u"毛南", u"撒拉", u"塔吉克", u"阿昌", u"普米", u"鄂温克",
                  u"怒", u"京", u"基诺", u"德昂", u"保安", u"俄罗斯", u"裕固", u"乌兹别克",
                  u"门巴", u"鄂伦春", u"独龙", u"塔塔尔", u"赫哲", u"珞巴", u"布朗"]

        if numb_text > 1:
            return random.sample(nation, numb_text)
        else:
            return random.choice(nation)

    @staticmethod
    def gen_pank_pan(name):
        pan = ""
        rule = [["AAA", "ZZZ"],
                ["P", "F", "C", "H", "A", "T", "B", "L", "J", "G"],
                ["-"], [1, 9999], ["A", "K"]]
        a = ord('A')
        a = random.randint(a, ord('Z'))

        pan += chr(a)
        a = random.randint(a, ord('Z'))

        pan += chr(a)
        a = random.randint(a, ord('Z'))

        pan += chr(a)

        pan += random.choice(rule[1])

        ispn = name.split()

        pan += ispn[2][0]

        pan += "%04d" % (random.randint(1, 9999))

        pan += chr(random.randint(ord('A'), ord('Z')))

        return pan

    @staticmethod
    def gen_pank_name():
        name = ""

        for i in range(3):
            num_alpha = random.randint(3, 7)

            for j in range(num_alpha):
                ch = random.randint(ord('A'), ord('Z'))
                name += chr(ch)

            if i != 2:
                name += " "

        if random.choice([True, False]):
            dot_num = random.randint(0, 3)

            n = len(name) - 2
            for i in range(dot_num):
                times = 10
                while times > 0:
                    pos = random.randint(2, n)
                    if name[pos] != '.' and name[pos] != ' ' and \
                            name[pos - 1] != '.' and name[pos - 1] != ' ' and \
                            name[pos + 1] != '.' and name[pos + 1] != ' ':
                        break

                    times -= 1
                if times != 0:
                    name = name[0:pos] + '.' + name[pos + 1:]

        return name
